import React, { createContext, useContext, useEffect, useState } from 'react';

interface Message {
  id: string;
  role: 'user' | 'assistant';
  content: string;
  timestamp: number;
}

interface ChatSession {
  id: string;
  title: string;
  lastUpdated: number;
  messages: Message[];
}

interface ChatContextType {
  currentSession: string | null;
  sessions: ChatSession[];
  messages: Message[];
  addMessage: (role: 'user' | 'assistant', content: string) => void;
  clearMessages: () => void;
  createNewSession: () => void;
  switchSession: (sessionId: string) => void;
  deleteSession: (sessionId: string) => void;
  updateSessionTitle: (sessionId: string, title: string) => void;
}

const ChatContext = createContext<ChatContextType | undefined>(undefined);

const generateId = () => Math.random().toString(36).substr(2, 9);

export const ChatProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
  // 从 localStorage 加载所有会话
  const [sessions, setSessions] = useState<ChatSession[]>(() => {
    const saved = localStorage.getItem('chatSessions');
    return saved ? JSON.parse(saved) : [];
  });

  // 当前会话ID
  const [currentSession, setCurrentSession] = useState<string | null>(() => {
    const saved = localStorage.getItem('currentSession');
    if (saved && sessions.some(session => session.id === saved)) {
      return saved;
    }
    return null;
  });

  // 获取当前会话的消息
  const messages = currentSession 
    ? sessions.find(s => s.id === currentSession)?.messages || []
    : [];

  // 保存会话到 localStorage
  useEffect(() => {
    localStorage.setItem('chatSessions', JSON.stringify(sessions));
  }, [sessions]);

  // 保存当前会话ID
  useEffect(() => {
    if (currentSession) {
      localStorage.setItem('currentSession', currentSession);
    } else {
      localStorage.removeItem('currentSession');
    }
  }, [currentSession]);

  const addMessage = (role: 'user' | 'assistant', content: string) => {
    const newMessage: Message = {
      id: generateId(),
      role,
      content,
      timestamp: Date.now(),
    };

    setSessions(prevSessions => {
      if (!currentSession) {
        // 如果没有当前会话，创建新会话
        const newSession: ChatSession = {
          id: generateId(),
          title: content.slice(0, 30) + '...',
          lastUpdated: Date.now(),
          messages: [newMessage],
        };
        setCurrentSession(newSession.id);
        return [...prevSessions, newSession];
      }

      // 更新现有会话
      return prevSessions.map(session => {
        if (session.id === currentSession) {
          return {
            ...session,
            lastUpdated: Date.now(),
            messages: [...session.messages, newMessage],
          };
        }
        return session;
      });
    });
  };

  const createNewSession = () => {
    const newSession: ChatSession = {
      id: generateId(),
      title: '新对话',
      lastUpdated: Date.now(),
      messages: [],
    };
    setSessions(prev => [...prev, newSession]);
    setCurrentSession(newSession.id);
  };

  const switchSession = (sessionId: string) => {
    if (sessions.some(session => session.id === sessionId)) {
      setCurrentSession(sessionId);
    }
  };

  const deleteSession = (sessionId: string) => {
    setSessions(prev => prev.filter(session => session.id !== sessionId));
    if (currentSession === sessionId) {
      const remainingSessions = sessions.filter(session => session.id !== sessionId);
      setCurrentSession(remainingSessions.length > 0 ? remainingSessions[0].id : null);
    }
  };

  const updateSessionTitle = (sessionId: string, title: string) => {
    setSessions(prev => prev.map(session => 
      session.id === sessionId ? { ...session, title } : session
    ));
  };

  const clearMessages = () => {
    if (currentSession) {
      setSessions(prev => prev.map(session => 
        session.id === currentSession ? { ...session, messages: [] } : session
      ));
    }
  };

  return (
    <ChatContext.Provider value={{
      currentSession,
      sessions,
      messages,
      addMessage,
      clearMessages,
      createNewSession,
      switchSession,
      deleteSession,
      updateSessionTitle,
    }}>
      {children}
    </ChatContext.Provider>
  );
};

export const useChat = () => {
  const context = useContext(ChatContext);
  if (context === undefined) {
    throw new Error('useChat must be used within a ChatProvider');
  }
  return context;
}; 